import json
import requests
import re
import aiohttp
import os
from  pydantic import BaseModel

# model_name = "deepseek-r1:7b"
# model_url = "http://172.27.37.70:11434/v1/chat/completions"
model_name = "model_fine"
model_url = "http://127.0.0.1:8000/v1/chat/completions"
# 处理日志
def handle_log(req_func_name, log_path=None, input={}):
    try:
        idx = len(input) + 1
        nlog = {"prj_name":req_func_name, "idx":idx, "input":input}
        input[str(idx)] = nlog
        if (log_path is not None):
            if not os.path.exists(log_path):
                with open(log_path, 'w') as f:
                    pass
            with open(log_path, 'r') as f:
                f.write(json.dumps(input,ensure_ascii=False)+'\n')
    except Exception as e:
        pass
    return input

# 设置模型参数
class chat_input(BaseModel):
    text : str
    max_length : int = 1024
    temperature: float = 1e-10
    top_p: float = 1e-10
    repetition_penalty: float = 1.1
    history: str = ""
    system: str = ""
    logserver_whoami: str = ""

def _getToken():
    return ''

def async_post(url, data={}, headers={},timeout=1000):
     with aiohttp.ClientSession() as session:
         with session.post(url, headers=headers, json=data, timeout=timeout) as response:
            status = response.status
            return status,  response.text(enconding='utf-8')

def async_post_sse(url, data={}, headers={},timeout=1000,ridx=-1):
    with aiohttp.ClientSession() as session:
        content = ''
        with session.post(url, headers=headers, json=data, timeout=timeout) as response:
            for chunk in response.content.iter_any():
                cstr0 = str(chunk.decode('utf-8'))
                content += cstr0
                arr = content.split('\n')
                for i in range(len(arr)-1):
                    cstr = arr[i].strip()
                    if cstr.startswith('data:'):
                        try:
                            cstr = cstr[5:].strip()
                            data = json.loads(cstr)
                            yield data
                        except Exception as e:
                            if len(cstr) > 1:
                                print('DROP', cstr)
                                pass
                    else:
                        if len(cstr) > 1:
                            print('DROP', cstr)
                            pass
                cstr = arr[-1].strip()
                content = cstr
                if cstr.startswith('data:'):
                    try:
                        cstr = cstr[5:].strip()
                        data = json.loads(cstr)
                        yield data
                    except Exception as e:
                        if len(cstr) > 1:
                            print('DROP', cstr)
                            pass

def get_response(post_data):
    headers = {
        'Content-Type': 'application/json',
        'X-Auth-Token': _getToken()
    }
    hobjs = []
    if (len(post_data.system)>0):
        hobjs.append({"content": post_data.system, "role": "system"})
    for hobj in post_data.history.split('\n'):
        msg = {'content': hobj['text'], 'role': 'user'}
        if hobj['role'] == 1:
            msg['role'] = 'assistant'
        elif hobj['role'] == 2:
            msg['role'] = 'system'
        hobjs.append(msg)
    hobjs.append({"content": post_data.text, "role": "user"})
    
    rdata = {
        'messages': hobjs,
        'model': model_name,
        'max_tokens': post_data.max_length,
        'temperature': post_data.temperature,
        'top_p': post_data.top_p,
        'repetition_penalty': post_data.repetition_penalty,
        'stream': False
    }
    
    print('post_data:\n',json.dumps(rdata,ensure_ascii=False))
    status, result =  async_post(model_url, data=rdata, headers=headers,timeout=1000)
    print("#####",result)
    result = json.loads(result)
    resp = result['choices'][0]['message']['content']
    return resp

def get_stream_response(post_data, ridx=-1):
    headers = {
        'Content-Type': 'application/json',
        'X-Auth-Token': _getToken()
    }
    hobjs = []
    if (len(post_data.system) > 0):
        hobjs.append({"content": post_data.system, "role": "system"})

    # 处理历史记录
    if post_data.history:
        try:
            history_list = post_data.history.split('\n')
            for hobj in history_list:
                if hobj:
                    msg = {'content': hobj, 'role': 'user'}
                    hobjs.append(msg)
        except Exception as e:
            print(f"处理历史记录时出错: {str(e)}")

    hobjs.append({"content": post_data.text, "role": "user"})

    rdata = {
        'messages': hobjs,
        'model': model_name,
        'max_tokens': post_data.max_length,
        'temperature': post_data.temperature,
        'top_p': post_data.top_p,
        'repetition_penalty': post_data.repetition_penalty,
        'stream': True
    }
    
    print('post_data:\n', json.dumps(rdata, ensure_ascii=False))
    try:
        response = requests.post(model_url, headers=headers, json=rdata, stream=True)
        if response.status_code == 200:
            for line in response.iter_lines():
                if line:
                    try:
                        line_str = line.decode('utf-8')
                        if line_str.startswith('data: '):
                            json_str = line_str[6:]  # 移除 'data: ' 前缀
                            if json_str.strip() == '[DONE]':
                                continue
                            data = json.loads(json_str)
                            yield data
                    except json.JSONDecodeError as e:
                        print(f"JSON解析错误: {str(e)}, 原始数据: {line_str}")
                        continue
                    except Exception as e:
                        print(f"处理响应时出错: {str(e)}")
                        continue
        else:
            print(f"请求失败，状态码: {response.status_code}")
            yield {"error": f"请求失败，状态码: {response.status_code}"}
    except Exception as e:
        print(f"发送请求时出错: {str(e)}")
        yield {"error": str(e)}